from __future__ import division

from base_inference_engines import BaseInferenceEngine
from result import Result
import paddle
from baidu_detector import Detector,load_predictor,create_inputs
import numpy as np

confidence_threshold = 0.5

class BaiduInferenceEngine(BaseInferenceEngine):
    def __init__(self):
        super().__init__()
        self.detector = None

    def load_model(self, model_path: str, **kwargs) -> bool:
        # Implement Baidu model loading using Paddle, etc.
        print("Loading model for Baidu...")
        paddle.enable_static()

        with open('infer_cfg.yml') as f:
            yml_conf = yaml.safe_load(f)
        arch = yml_conf['arch']
        detector_func = 'Detector'

        self.detector = eval(detector_func)(
            model_path)

        return True

    def run_inference(self, input_image: np.ndarray, **kwargs) -> dict:
        # Implement Baidu-specific inference, e.g., using Paddle
        print("Running inference on Baidu...")

        #img_list = get_test_images(None,'test_10240_v2.png')
        img_list = ['test_10240_v2.png']
        cregions = self.detector.predict_image_slice(img_list)
        # Format results to match the unified output structure
        print(cregions)
        # return {"boxes": [], "scores": [], "labels": []}
        boxes_value = cregions.get("boxes")
        result_list = []
        for index, box in enumerate(boxes_value):
            if box[4] < confidence_threshold:
                continue
            bbox = Result.BBox(box[2], box[3], box[4], box[5])
            result = Result(bbox=bbox, confidence=box[1], label=box[0])
            result_list.append(result)
        # Format results to match the unified output structure
        return result_list




